#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import print_function
import rospy
import math
import tf
from tf import transformations
from sensor_msgs.msg import LaserScan
from nav_msgs.msg import Odometry
from six.moves import range

def deg2rad(x):
    return x * math.pi / 180

def rad2deg(x):
    return x * 180 / math.pi

class Sensor(object):
    def __init__(self, median_filter_window = 3, simulation = "yes"):
        self.sub_leader_odom = rospy.Subscriber('/tb3_0/odom', Odometry, self.leader_odom_CB, queue_size=1,buff_size=52428800)
        self.sub_follower1_odom = rospy.Subscriber('/tb3_1/odom', Odometry, self.follower1_odom_CB, queue_size=1,buff_size=52428800)
        self.sub_follower2_odom = rospy.Subscriber('/tb3_2/odom', Odometry, self.follower2_odom_CB, queue_size=1,buff_size=52428800)

        self.sub_raw_leader_scan = rospy.Subscriber('/tb3_0/scan', LaserScan, self.raw_leader_scan_CB, queue_size=1,buff_size=52428800) 
        self.sub_raw_follower1_scan = rospy.Subscriber('/tb3_1/scan', LaserScan, self.raw_follower1_scan_CB, queue_size=1,buff_size=52428800)
        self.sub_raw_follower2_scan = rospy.Subscriber('/tb3_2/scan', LaserScan, self.raw_follower2_scan_CB, queue_size=1,buff_size=52428800)

        self.pub_leader_scan = rospy.Publisher('/tb3_0/processed_scan', LaserScan, queue_size=1)
        self.pub_follower1_scan = rospy.Publisher('/tb3_1/processed_scan', LaserScan, queue_size=1)
        self.pub_follower2_scan = rospy.Publisher('/tb3_2/processed_scan', LaserScan, queue_size=1)

        self.simulation = simulation

        self.range_min = 0.119999997318
        if (simulation == "yes"):
            self.range_max = 3.5 
        else:
            self.range_max = 3.5
        self.length_threshold = 0.5 # 设置感知距离
        self.hidden_robots_length_threshold = 0.5
        self.median_filter_window = median_filter_window
        
        # leader
        self.leader_ranges = []
        self.cur_leader_data = LaserScan()
        self.processed_leader_data = LaserScan()
        self.window_list = []

        # follower1
        self.follower1_ranges = []
        self.cur_follower1_data = LaserScan()
        self.processed_follower1_data = LaserScan()
        self.window_list_follower1 = []

        # follower2
        self.follower2_ranges = []
        self.cur_follower2_data = LaserScan()
        self.processed_follower2_data = LaserScan()
        self.window_list_follower2 = []
    
    def run(self):
        if ((len(self.cur_leader_data.ranges) == 360) and (len(self.cur_follower1_data.ranges) == 360) and \
             (len(self.cur_follower2_data.ranges) == 360)):
            self.preprocess()
            if (self.simulation == "yes"):
                self.median_filter()
                self.hidden_robots()
            self.publish()

    def hidden_robots(self):
        self.leader_hidden_robots()
        self.follower1_hidden_robots()
        self.follower2_hidden_robots()

    def leader_hidden_robots(self):
        print("leader_hidden_robots()")
        for i in range(len(self.leader_ranges)):
            if (self.leader_ranges[i] >= self.range_min)and(self.leader_ranges[i] <= self.range_max):
                theta = deg2rad(i)
                distance = self.leader_ranges[i]
                sensor_x = self.leader_x + distance * math.cos(theta + self.leader_yaw)
                sensor_y = self.leader_y + distance * math.sin(theta + self.leader_yaw) 
                if (self.distance(sensor_x, sensor_y, self.follower1_x, self.follower1_y) < self.hidden_robots_length_threshold or \
                    self.distance(sensor_x, sensor_y, self.follower2_x, self.follower2_y) < self.hidden_robots_length_threshold):
                    self.leader_ranges[i] = 3.5
    
    def follower1_hidden_robots(self):
        print("follower1_hidden_robots()")
        for i in range(len(self.follower1_ranges)):
            if (self.follower1_ranges[i] >= self.range_min)and(self.follower1_ranges[i] <= self.range_max):
                theta = deg2rad(i)
                distance = self.follower1_ranges[i]
                sensor_x = self.follower1_x + distance * math.cos(theta + self.follower1_yaw)
                sensor_y = self.follower1_y + distance * math.sin(theta + self.follower1_yaw) 
                if (self.distance(sensor_x, sensor_y, self.leader_x, self.leader_y) < self.hidden_robots_length_threshold or \
                    self.distance(sensor_x, sensor_y, self.follower2_x, self.follower2_y) < self.hidden_robots_length_threshold):
                    self.follower1_ranges[i] = 3.5
                
    def follower2_hidden_robots(self):
        print("follower2_hidden_robots()")
        for i in range(len(self.follower2_ranges)):
            if (self.follower2_ranges[i] >= self.range_min)and(self.follower2_ranges[i] <= self.range_max):
                theta = deg2rad(i)
                distance = self.follower2_ranges[i]
                sensor_x = self.follower2_x + distance * math.cos(theta + self.follower2_yaw)
                sensor_y = self.follower2_y + distance * math.sin(theta + self.follower2_yaw) 
                if (self.distance(sensor_x, sensor_y, self.leader_x, self.leader_y) < self.hidden_robots_length_threshold or \
                    self.distance(sensor_x, sensor_y, self.follower1_x, self.follower1_y) < self.hidden_robots_length_threshold):
                    self.follower2_ranges[i] = 3.5

    def distance(self, x1, y1, x2, y2):
        return math.sqrt(math.pow(x1 - x2, 2) + math.pow(y1 - y2, 2))
        
    def publish(self):
        # print("publish()")
        # leader
        self.processed_leader_data.ranges = self.leader_ranges
        self.processed_leader_data.header.frame_id = "/tb3_0/base_footprint"
        self.processed_leader_data.header.stamp = rospy.get_rostime()
        self.pub_leader_scan.publish(self.processed_leader_data)
        # follower1
        self.processed_follower1_data.ranges = self.follower1_ranges
        self.processed_follower1_data.header.frame_id = "/tb3_1/base_footprint"
        self.processed_follower1_data.header.stamp = rospy.get_rostime()
        self.pub_follower1_scan.publish(self.processed_follower1_data)
        # follower2
        self.processed_follower2_data.ranges = self.follower2_ranges
        self.processed_follower2_data.header.frame_id = "/tb3_2/base_footprint"
        self.processed_follower2_data.header.stamp = rospy.get_rostime()
        self.pub_follower2_scan.publish(self.processed_follower2_data)


    # 处理实际雷达数据不在限制区域内的情况(只考虑0.12 ～ 0.5m 内的障碍物)
    def preprocess(self):
        self.processed_leader_data = self.cur_leader_data
        self.processed_follower1_data = self.cur_follower1_data
        self.processed_follower2_data = self.cur_follower2_data
        self.leader_ranges = list(self.processed_leader_data.ranges)
        self.follower1_ranges = list(self.processed_follower1_data.ranges)
        self.follower2_ranges = list(self.processed_follower2_data.ranges)

        if (len(self.leader_ranges) == 360) and (len(self.follower1_ranges) == 360) and (len(self.follower2_ranges) == 360):
            # print("preprocess()")
            for i in range(360):
                if (self.leader_ranges[i] < self.range_min)or(self.leader_ranges[i] > self.range_max):
                    self.leader_ranges[i] = 3.5
                if (self.follower1_ranges[i] < self.range_min)or(self.follower1_ranges[i] > self.range_max):
                    self.follower1_ranges[i] = 3.5               
                if (self.follower2_ranges[i] < self.range_min)or(self.follower2_ranges[i] > self.range_max):
                    self.follower2_ranges[i] = 3.5

    # 中值滤波
    def median_filter(self):
        print("median_filter()")
        for i in range(len(self.leader_ranges)):
            if len(self.window_list) < self.median_filter_window:
                self.window_list.append(self.leader_ranges[i])
                self.window_list_follower1.append(self.follower1_ranges[i])
                self.window_list_follower2.append(self.follower2_ranges[i])
            else:
                # print("111111111111111111111111111111111111111111111")
                # print(self.ranges[i])
                self.window_list.append(self.leader_ranges[i])
                self.window_list_follower1.append(self.follower1_ranges[i])
                self.window_list_follower2.append(self.follower2_ranges[i])
                # print(self.window_list)
                self.window_list.pop(0)
                self.window_list_follower1.pop(0)
                self.window_list_follower2.pop(0)
                # print(self.window_list)
                leader_sorted_list = sorted(self.window_list)
                follower1_sorted_list = sorted(self.window_list_follower1)
                follower2_sorted_list = sorted(self.window_list_follower2)
                # print(sorted_list)
                # self.leader_ranges[i] = leader_sorted_list[self.median_filter_window / 2]
                self.leader_ranges[i] = leader_sorted_list[int(self.median_filter_window / 2)]
                # self.follower1_ranges[i] = follower1_sorted_list[self.median_filter_window / 2]
                self.follower1_ranges[i] = follower1_sorted_list[int(self.median_filter_window / 2)]
                # self.follower2_ranges[i] = follower2_sorted_list[self.median_filter_window / 2]
                self.follower2_ranges[i] = follower2_sorted_list[int(self.median_filter_window / 2)]
                # print(self.ranges[i])
                # print("22222222222222222222222222222222222222222222222")
                if (i == 360):
                    self.window_list = []
                    self.window_list_follower1 = []
                    self.window_list_follower2 = []
            

    # 领航者回调函数
    def raw_leader_scan_CB(self, data):
        self.cur_leader_data = data
        print(("leader range len : %d" %len(data.ranges)))
        # self.processed_leader_data = self.cur_leader_data
        # self.ranges = list(data.ranges)
        # print(len(self.ranges))
        # print(type(self.ranges))

    # follower1
    def raw_follower1_scan_CB(self, data):
        self.cur_follower1_data = data
        print(("follower1 range len : %d" %len(data.ranges)))

    # follower2
    def raw_follower2_scan_CB(self, data):
        self.cur_follower2_data = data
        print(("follower2 range len : %d" %len(data.ranges)))
    
    def leader_odom_CB(self, data):
        self.leader_x = data.pose.pose.position.x
        self.leader_y = data.pose.pose.position.y
        a = data.pose.pose.orientation.x
        b = data.pose.pose.orientation.y
        c = data.pose.pose.orientation.z
        d = data.pose.pose.orientation.w
        (roll,pitch,yaw) = transformations.euler_from_quaternion([a, b, c, d])
        self.leader_yaw = yaw
        # print("ledaer_x: %lf, ledaer_y: %lf" %(self.ledaer_x, self.ledaer_y))

    def follower1_odom_CB(self, data):
        self.follower1_x = data.pose.pose.position.x
        self.follower1_y = data.pose.pose.position.y
        a = data.pose.pose.orientation.x
        b = data.pose.pose.orientation.y
        c = data.pose.pose.orientation.z
        d = data.pose.pose.orientation.w
        (roll,pitch,yaw) = transformations.euler_from_quaternion([a, b, c, d])
        self.follower1_yaw = yaw
        # print("follower1_x: %lf, follower1_y: %lf" %(self.follower1_x, self.follower1_y))

    def follower2_odom_CB(self, data):
        self.follower2_x = data.pose.pose.position.x
        self.follower2_y = data.pose.pose.position.y
        a = data.pose.pose.orientation.x
        b = data.pose.pose.orientation.y
        c = data.pose.pose.orientation.z
        d = data.pose.pose.orientation.w
        (roll,pitch,yaw) = transformations.euler_from_quaternion([a, b, c, d])
        self.follower2_yaw = yaw
        # print("follower2_x: %lf, follower2_y: %lf" %(self.follower2_x, self.follower2_y))

def node():
    rospy.init_node('sensor_node')
    rospy.loginfo("sensor node is started!!")
    simulation = rospy.get_param('~simulation')
    sensor = Sensor(simulation = simulation)
    rate = rospy.Rate(5)
    while not rospy.is_shutdown():
        sensor.run()
        rate.sleep()

if __name__ == '__main__':
    try:
        node()
    except rospy.ROSInterruptException:
        pass 
